"""
文件名称: system_initialization.py
功能: 系统初始化函数文件：随机基站用户位置、随机信道状态等
作者: 喻越
创建日期: 2025-01-15
"""
import numpy as np
import matplotlib.pyplot as plt
from ..parameter.constant import const
import cvxpy as cp


def uniform2d(num, L):
    '''
    均匀分布
    :param num: 均匀分布点数
    :param L:仿真区域边长
    :return: 均匀分布点的二维坐标
    '''
    position = np.random.uniform(0, 1, [num, 2])
    position = position * L
    return position


def ppp2d(num, L):
    '''
    泊松点分布
    :param num: 泊松点分布参数
    :param L:仿真区域边长
    :return: 泊松分布点的二维坐标
    '''
    num_actual = np.random.poisson(num)
    position = np.random.uniform(0, 1, [num_actual, 2])
    position = position * L
    return position


def pcp2d(cluster_center, num_point, r_max):
    '''
    泊松簇分布
    :param cluster_center:簇中心二维坐标
    :param num_point:每簇的点数
    :param r_max:最大簇半径
    :return:簇的点位置
    '''
    num_center = cluster_center.shape[0]  # 簇中心数
    theta = 2 * np.pi * np.random.rand(num_point, num_point)  # 簇点角度
    cluster_r = r_max * np.random.rand(num_point, num_point)  # 簇点半径
    position = np.empty(shape=[num_center * num_point, 2])  # 点位置
    for i in range(num_center):
        for j in range(num_point):
            position[i * j, 0] = cluster_center[i, 0] + cluster_r[i, j] * np.cos(theta[i, j])
            position[i * j, 1] = cluster_center[i, 1] + cluster_r[i, j] * np.sin(theta[i, j])
    return position


def macro_bs_init(bs_pos, L):
    '''
    在区域中心位置放置宏基站
    :param bs_pos: 微基站二维位置
    :param L: 仿真区域边长
    :return: 所有基站的二维位置
    '''
    macro_bs = np.array([[L / 2, L / 2]])
    bs_pos = np.vstack([macro_bs, bs_pos])
    return bs_pos


def distance_calculation(bs_pos, ue_pos):
    '''
    计算各基站与各用户间距离
    :param bs_pos: 基站二维坐标
    :param ue_pos: 用户二维坐标
    :return: 基站用户间距离
    '''
    bs_num = bs_pos.shape[0]
    ue_num = ue_pos.shape[0]
    l_bs_ue = np.empty(shape=[ue_num, bs_num])  # 用户与基站的距离矩阵
    for j in range(int(ue_num)):
        for i in range(int(bs_num)):
            l_bs_ue[j, i] = np.sqrt(np.sum(np.square(bs_pos[i, :] - ue_pos[j, :])))
    return l_bs_ue


def user_adjust_overloaded(x_bs_ue, bs_status, rb_req):
    '''
    微基站过载时，转移资源块需求最多的用户到其他没有超载的基站，如果基站全部超载，x全部置-1，当前操作违规
    :param x_bs_ue: 基站用户连接  # shape = [ue_num,bs_num]
    :param bs_status: 基站休眠状态  # shape = [1,bs_num]
    :param rb_req: 用户需求状态  # shape = [ue_num,bs_num]
    :return: x_bs_ue: 调整后的基站用户连接  # shape = [ue_num,bs_num]
    '''
    bs_num = bs_status.shape[1]
    bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
    rb_max = np.hstack((np.full([1, 1], const.RB_MAX_MACRO), np.full([1, bs_num - 1], const.RB_MAX_MICRO)))
    # 超过100轮，说明有用户接入哪个基站都会造成某个基站超载，其他基站不超载
    k = 0
    loop = 40
    while ((bs_rb_access - rb_max) >= 1e-5).any():
        print('用户调整第' + str(k) + '轮！')
        out_bs_index = np.where(bs_rb_access[0] > rb_max[0])  # 超载基站
        less_bs_index = np.intersect1d(np.where(bs_rb_access[0] < rb_max[0]), np.where(bs_status[0] == 1))  # 工作中的未超载基站
        if len(less_bs_index) == 0 or k >= loop:  # 工作的全部超载
            x_bs_ue[:, :] = -1
            break
        less_bs_index = tuple([less_bs_index])
        # exchange_ue_index = np.where(x_bs_ue[:,out_bs_index]==1)
        # # print(exchange_ue_index[0])
        j = np.random.choice(out_bs_index[0])  # 超载基站中随机选

        j_next = np.random.choice(less_bs_index[0])  # 未超载基站中随机选
        i = np.argwhere((x_bs_ue * rb_req)[:, j] == np.max((x_bs_ue * rb_req)[:, j]))[0][0]  # 随机选取需求最大的用户接入宏基站
        x_bs_ue[i, j] = 0
        x_bs_ue[i, j_next] = 1  # 随机选取用户接入未超载基站
        bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
        k = k + 1
        # if ((bs_rb_access-rb_max) >= 1e-5).all():
        #     x_bs_ue[:, :] = -1
        #     break
    return x_bs_ue


def user_rejected_overloaded(x_bs_ue, bs_status, rb_req):
    '''
    微基站过载时，基站拒绝资源块需求最多的用户
    :param x_bs_ue: 基站用户连接  # shape = [ue_num,bs_num]
    :param bs_status: 基站休眠状态  # shape = [1,bs_num]
    :param rb_req: 用户需求状态  # shape = [ue_num,bs_num]
    :return: x_bs_ue: 调整后的基站用户连接  # shape = [ue_num,bs_num]
    '''
    bs_num = bs_status.shape[1]
    bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
    rb_max = np.hstack((np.full([1, 1], const.RB_MAX_MACRO), np.full([1, bs_num - 1], const.RB_MAX_MICRO)))
    # 超过100轮，说明有用户接入哪个基站都会造成某个基站超载，其他基站不超载
    k = 0
    loop = 100
    while ((bs_rb_access - rb_max) >= 1e-5).any():
        print('用户调整第' + str(k) + '轮！')
        out_bs_index = np.where(bs_rb_access[0] > rb_max[0])  # 超载基站
        less_bs_index = np.intersect1d(np.where(bs_rb_access[0] < rb_max[0]), np.where(bs_status[0] == 1))  # 工作中的未超载基站
        if len(out_bs_index) == 0 or k >= loop:  # 如果没有超载基站
            break
        less_bs_index = tuple([less_bs_index])
        # exchange_ue_index = np.where(x_bs_ue[:,out_bs_index]==1)
        # # print(exchange_ue_index[0])
        j = np.random.choice(out_bs_index[0])  # 超载基站中随机选
        i = np.argwhere((x_bs_ue * rb_req)[:, j] == np.max((x_bs_ue * rb_req)[:, j]))[0][0]  # 随机选取需求最大的用户拒绝
        x_bs_ue[i, j] = 0
        bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
        k = k + 1
        # if ((bs_rb_access-rb_max) >= 1e-5).all():
        #     x_bs_ue[:, :] = -1
        #     break
    return x_bs_ue


def nearby_access(bs_pos, ue_pos, bs_status, rb_req):
    '''
    用户就近接入工作基站;如果微基站超载，转移到其他没有超载的基站，如果基站全部超载，拒绝需求最大的用户
    :param bs_pos: 基站二维坐标
    :param ue_pos: 用户二维坐标
    :param bs_status: 基站休眠状态
    :param rb_req：需求资源块
    :return: 用户就近接入基站的连接矩阵
    '''
    bs_num = bs_pos.shape[0]
    ue_num = ue_pos.shape[0]
    l_bs_ue = distance_calculation(bs_pos, ue_pos)  # shape=[ue_num, bs_num]
    for i in range(bs_num):
        if bs_status[0][i] == 0:
            l_bs_ue[:, i] = -1  # 休眠基站失效

    x_bs_ue = np.zeros(shape=[ue_num, bs_num])  # 基站用户间连接矩阵

    for j in range(int(ue_num)):
        # 查找大于0的列表中的最小值
        l_ue = l_bs_ue[j, :]
        l_ue_min = min(k for k in l_ue if k > 0)
        i = np.where(l_ue == l_ue_min)
        x_bs_ue[j, i] = int(1)

    # 微基站过载时，转移到其他没有超载的基站，如果基站全部超载，x全部置-1，当前操作违规
    # x_bs_ue = user_adjust_overloaded(x_bs_ue, bs_status, rb_req)
    # 微基站过载时，拒绝需求最大的用户
    x_bs_ue = user_rejected_overloaded(x_bs_ue, bs_status, rb_req)

    return x_bs_ue


def ratemax_access(bs_pos, ue_pos, bs_status, rb_req, l_bs_ue, h_bs_ue):
    '''
    用户接入提供最大速率的工作基站
    '''
    bs_num = bs_pos.shape[0]
    ue_num = ue_pos.shape[0]
    one = np.ones(shape=[ue_num, bs_num])
    RB_W = np.tile(const.RB_W, (ue_num, 1))
    x_bs_ue = np.ones(shape=[ue_num, bs_num])
    # 用户速率矩阵 shape=[ue_num, bs_num]
    r_ue_temp = RB_W * (x_bs_ue * rb_req) * np.log2(
        one + (const.RB_PR * rb_req * h_bs_ue) / (np.power(l_bs_ue, const.ALPHA) * const.NOISE))
    order_bs_ue = np.argsort(r_ue_temp, axis=1)
    order_ue = np.zeros(shape=[ue_num, 1])  # 每个用户的志愿顺序0 1 2 3
    x_bs_ue = (r_ue_temp == np.max(r_ue_temp, axis=1, keepdims=True)).astype(int)
    # 基站过载时，需求最大的用户被拒绝，进入第二志愿
    bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
    rb_max = np.hstack((np.full([1, 1], const.RB_MAX_MACRO), np.full([1, bs_num - 1], const.RB_MAX_MICRO)))
    k = 0
    loop = 100
    while ((bs_rb_access - rb_max) >= 1e-5).any():
        print('用户调整第' + str(k) + '轮！')
        out_bs_index = np.where(bs_rb_access[0] > rb_max[0])  # 超载基站
        less_bs_index = np.intersect1d(np.where(bs_rb_access[0] < rb_max[0]), np.where(bs_status[0] == 1))  # 工作中的未超载基站
        if len(out_bs_index) == 0 or k >= loop:  # 如果没有超载基站
            break
        less_bs_index = tuple([less_bs_index])
        j = np.random.choice(out_bs_index[0])  # 超载基站中随机选
        i = np.argwhere((x_bs_ue * rb_req)[:, j] == np.max((x_bs_ue * rb_req)[:, j]))[0][0]  # 随机选取需求最大的用户拒绝
        x_bs_ue[i, j] = 0
        # 选择第二速率
        if order_ue[i] < bs_num:  # 志愿已用完
            index = np.where(order_bs_ue[i] == int(order_ue[i][0] + 1))[0]
            x_bs_ue[i, index] = 1
            order_ue[i][0] = order_ue[i][0] + 1
        bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
        k = k + 1
    return x_bs_ue


def ratemax_access_rejected(bs_pos, ue_pos, bs_status, rb_req, l_bs_ue, h_bs_ue):
    '''
    用户接入提供最大速率的工作基站。超载基站拒绝最大需求的用户
    '''
    bs_num = bs_pos.shape[0]
    ue_num = ue_pos.shape[0]
    one = np.ones(shape=[ue_num, bs_num])
    RB_W = np.tile(const.RB_W, (ue_num, 1))
    x_bs_ue = np.ones(shape=[ue_num, bs_num])
    # 用户速率矩阵 shape=[ue_num, bs_num]
    r_ue_temp = RB_W * (x_bs_ue * rb_req) * np.log2(
        one + (const.RB_PR * rb_req * h_bs_ue) / (np.power(l_bs_ue, const.ALPHA) * const.NOISE))
    x_bs_ue = (r_ue_temp == np.max(r_ue_temp, axis=1, keepdims=True)).astype(int)
    # 微基站过载时，拒绝需求最大的用户
    x_bs_ue = user_rejected_overloaded(x_bs_ue, bs_status, rb_req)
    return x_bs_ue


def random_access(bs_num, ue_num, bs_status, rb_req):
    '''
    用户随机接入工作基站
    :param bs_num: 基站数量
    :param ue_num: 用户数量
    :param bs_status: 基站休眠状态
    :return: 用户基站连接矩阵
    '''
    x_bs_ue = np.zeros(shape=[ue_num, bs_num])  # 基站用户间连接矩阵
    bs_all_index = np.array([i for i in range(bs_num)])  # shape=[bs_num,-] TODO
    bs_sleep_index = np.where(bs_status == 0)[1]  # shape=[N_bs_sleep,-]
    bs_work_index = np.isin(bs_all_index, bs_sleep_index, invert=True)  # 求补集 Bool值
    bs_work_index = bs_all_index[bs_work_index]  # 工作基站索引
    for j in range(int(ue_num)):
        i = np.random.choice(bs_work_index)
        x_bs_ue[j, i] = int(1)

    # 微基站过载时，转移到其他没有超载的基站，如果基站全部超载，x全部置-1，当前操作违规
    # x_bs_ue = user_adjust_overloaded(x_bs_ue, bs_status, rb_req)
    # 微基站过载时，拒绝需求最大的基站
    x_bs_ue = user_rejected_overloaded(x_bs_ue, bs_status, rb_req)
    return x_bs_ue


def channel_status(bs_num, ue_num):
    '''
    基站用户的信道状态
    :param bs_num: 基站数目
    :param ue_num: 用户数目
    :return: 用户基站信道矩阵
    '''
    h_bs_ue = np.random.exponential(1, size=[ue_num, bs_num])
    return h_bs_ue


def rb_demand(ue_snr_theta, l_bs_ue, h_bs_ue):
    '''
    求取用户连接到各个基站的需求功率、资源块
    :param ue_snr_theta: 用户SNR阈值
    :param bs_num: 基站数目
    :param ue_num: 用户数目
    :param bs_pos: 基站位置
    :param ue_pos: 用户位置
    :return: 用户连接到各个基站的需求功率、资源块
    '''
    ue_num = h_bs_ue.shape[0]
    bs_num = h_bs_ue.shape[1]
    ue_snr_theta_temp = np.zeros(shape=[ue_num, bs_num])  # 扩充snr阈值矩阵，便于计算
    for j in range(ue_num):
        ue_snr_theta_temp[j, :] = ue_snr_theta[j, 0]

    pr_req = const.NOISE * ue_snr_theta_temp * np.power(l_bs_ue,
                                                        const.ALPHA) / h_bs_ue  # 用户连接到各个基站的需求功率  shape=[ue_num, bs_num]
    rb_req = np.ceil(pr_req / const.RB_PR)  # 向上取整，用户连接到各个基站的需求资源块

    return pr_req, rb_req


def rate_sum_calculation(x_bs_ue, h_bs_ue, l_bs_ue, rb_req):
    '''
    计算各用户速率与总速率
    :param x_bs_ue: 连接矩阵
    :param h_bs_ue: 信道矩阵
    :param l_bs_ue: 距离矩阵
    :param rb_req: 需求资源块矩阵
    :return: 各用户在当前连接下的速率与总速率
    '''
    ue_num = x_bs_ue.shape[0]
    bs_num = x_bs_ue.shape[1]
    # 计算用户总速率
    one = np.ones(shape=[ue_num, bs_num])
    RB_W = np.tile(const.RB_W, (ue_num, 1))
    # 用户速率矩阵 shape=[ue_num, bs_num]
    r_ue_temp = RB_W * (x_bs_ue * rb_req) * np.log2(
        one + (const.RB_PR * rb_req * h_bs_ue) / (np.power(l_bs_ue, const.ALPHA) * const.NOISE))
    # 用户速率矩阵 shape=[ue_num, 1] 升维后转置
    r_ue = np.transpose(np.expand_dims(np.sum(r_ue_temp, axis=1), 0))
    rb_req_access = np.transpose(np.expand_dims(np.sum(x_bs_ue * rb_req, axis=1), 0))  # 用户在当前接入下所需的资源块
    # 用户服务质量满足情况 考虑用户未接入的情况
    ue_status = np.sum(x_bs_ue, 1)[:, None]
    delta_connected = (np.sum(ue_status)) / ue_num

    return r_ue, ue_status, delta_connected


def bs_energy_consumption_calculation(bs_status, x_bs_ue, rb_req):
    '''
    依据休眠情况和连接情况，计算基站能耗
    :param bs_status: 基站休眠状态
    :param x_bs_ue: 连接状态
    :param rb_req: 资源块需求
    :return: 基站能耗，基站已用资源块
    '''
    bs_num = x_bs_ue.shape[1]  # 求取基站数
    # 当前接入下各基站使用的资源块 shape=[1,bs_num]
    bs_rb_access = np.sum((x_bs_ue * rb_req), 0) * bs_status
    p_static = np.zeros(shape=[1, bs_num])  # 基站静态功率
    for i in range(bs_num):
        if i == 0:
            p_static[0][i] = const.P_STATIC_SLEEP_MACRO if bs_status[0][i] == 0 else const.P_STATIC_WORK_MACRO
        else:
            p_static[0][i] = const.P_STATIC_SLEEP_MICRO if bs_status[0][i] == 0 else const.P_STATIC_WORK_MICRO

    p_bs = const.RB_PR * bs_rb_access / const.ETA + p_static  # 包括休眠基站 shape=[1,bs_num]

    return p_bs, bs_rb_access


def p_pv_initialization(bs_num):
    '''
    光伏风力的仿真分布
    :param bs_num:基站数目
    :return:
    '''
    p_pv = np.zeros([1, bs_num])
    for i in range(bs_num):
        # pv_day = (np.random.poisson(const.PV_MAX)
        #           + np.random.poisson(const.WIND_MAX))
        pv_day = np.random.poisson(const.PV_MAX)
        if i == 0:
            p_pv[0][i] = pv_day * const.S_PV_MACRO * const.GAMMA_PV
        else:
            p_pv[0][i] = pv_day * const.S_PV_MICRO * const.GAMMA_PV
    return p_pv
